"""
知识库模块 - 支持问答对的知识库管理，集成向量数据库功能
"""
import json
import os
from typing import List, Tuple, Optional, Dict, Any
from .vector_db import KnowledgeVectorDB


class KnowledgeBase:
    """知识库类"""
    
    def __init__(self, file_path: str = None, dir_path: str = None):
        """
        初始化知识库
        
        Args:
            file_path: 知识库文件路径
            dir_path: 知识库目录路径（加载目录下所有知识库文件）
        """
        self.file_path = file_path
        self.dir_path = dir_path
        self.data = None
        self.knowledge_bases = {}  # 当使用dir_path时，存储多个知识库
        
        # 初始化向量数据库
        if dir_path:
            self.vector_db = KnowledgeVectorDB(dir_path)
        else:
            self.vector_db = None
        
        if file_path:
            self.load_knowledge_base(file_path)
        elif dir_path:
            self.load_knowledge_base_from_dir(dir_path)
    
    def load_knowledge_base(self, file_path: str = None) -> bool:
        """
        加载知识库文件
        
        Args:
            file_path: 知识库文件路径，如果为None则使用初始化时的file_path
            
        Returns:
            bool: 是否成功加载
        """
        if file_path is None:
            file_path = self.file_path
        
        if file_path is None:
            return False
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 处理新旧格式兼容性
            if isinstance(data, list):
                # 旧格式：直接是问答对列表
                self.data = {
                    'name': os.path.basename(file_path).replace('.json', ''),
                    'description': '自动转换的旧格式知识库',
                    'qa_pairs': data
                }
            else:
                # 新格式：包含元数据的字典
                self.data = data
                
                # 确保有qa_pairs字段
                if 'qa_pairs' not in self.data:
                    self.data['qa_pairs'] = []
            
            self.file_path = file_path
            return True
            
        except Exception as e:
            print(f"加载知识库文件 {file_path} 失败: {e}")
            return False
    
    def load_knowledge_base_from_dir(self, dir_path: str = None) -> bool:
        """
        从目录加载所有知识库文件
        
        Args:
            dir_path: 知识库目录路径，如果为None则使用初始化时的dir_path
            
        Returns:
            bool: 是否成功加载
        """
        if dir_path is None:
            dir_path = self.dir_path
        
        if dir_path is None:
            return False
        
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
            return True
        
        self.knowledge_bases = {}
        for filename in os.listdir(dir_path):
            if filename.endswith('.json'):
                kb_path = os.path.join(dir_path, filename)
                kb_name = filename.replace('.json', '')
                
                kb = KnowledgeBase(kb_path)
                if kb.data:
                    self.knowledge_bases[kb_name] = kb
        
        return True
    
    def match_question(self, question: str) -> Optional[Tuple[str, List[str]]]:
        """
        匹配问题，返回最匹配的问题和答案列表
        优先使用向量数据库进行语义搜索，如果失败则回退到关键词匹配
        
        Args:
            question: 用户问题
            
        Returns:
            Optional[Tuple[str, List[str]]]: (匹配的问题, 答案列表) 或 None
        """
        # 优先使用向量数据库进行语义搜索
        semantic_match = self._semantic_match(question)
        if semantic_match:
            return semantic_match
        
        # 如果语义搜索没有结果，回退到关键词匹配
        if self.data and 'qa_pairs' in self.data:
            return self._match_in_single_kb(question)
        elif self.knowledge_bases:
            return self._match_in_multiple_kbs(question)
        else:
            return None
    
    def _semantic_match(self, question: str) -> Optional[Tuple[str, List[str]]]:
        """使用向量数据库进行语义匹配"""
        if not self.vector_db:
            return None
        
        # 获取当前角色（从文件路径推断）
        role = None
        if self.file_path:
            role = os.path.basename(self.file_path).replace('.json', '')
        
        if not role and self.dir_path and self.knowledge_bases:
            # 如果是多知识库模式，尝试在所有角色中搜索
            for role_name in self.vector_db.get_all_roles():
                results = self.vector_db.search(role_name, question, top_k=1, threshold=0.25)
                if results:
                    doc, score, metadata = results[0]
                    if metadata.get('type') == 'question':
                        return (doc, [metadata['answer']])
                    elif metadata.get('type') == 'answer':
                        return (metadata.get('question', '相关问题'), [doc])
        elif role:
            # 单知识库模式
            results = self.vector_db.search(role, question, top_k=1, threshold=0.25)
            if results:
                doc, score, metadata = results[0]
                if metadata.get('type') == 'question':
                    return (doc, [metadata['answer']])
                elif metadata.get('type') == 'answer':
                    return (metadata.get('question', '相关问题'), [doc])
        
        return None
    
    def _match_in_single_kb(self, question: str) -> Optional[Tuple[str, List[str]]]:
        """在单个知识库中匹配问题"""
        best_match = None
        best_score = 0
        
        for qa in self.data['qa_pairs']:
            q = qa.get('question', '')
            a = qa.get('answer', '')
            
            if not q or not a:
                continue
            
            # 简单匹配算法：完全匹配或部分匹配
            score = self._calculate_match_score(question, q)
            
            if score > best_score:
                best_score = score
                best_match = (q, [a])
        
        # 设置匹配阈值
        if best_score >= 0.7:  # 70%匹配度
            return best_match
        
        return None
    
    def _match_in_multiple_kbs(self, question: str) -> Optional[Tuple[str, List[str]]]:
        """在多个知识库中匹配问题"""
        best_match = None
        best_score = 0
        best_kb_name = None
        
        for kb_name, kb in self.knowledge_bases.items():
            if kb.data and 'qa_pairs' in kb.data:
                for qa in kb.data['qa_pairs']:
                    q = qa.get('question', '')
                    a = qa.get('answer', '')
                    
                    if not q or not a:
                        continue
                    
                    score = self._calculate_match_score(question, q)
                    
                    if score > best_score:
                        best_score = score
                        best_match = (q, [a])
                        best_kb_name = kb_name
        
        if best_score >= 0.7:
            return best_match
        
        return None
    
    def _calculate_match_score(self, question: str, target: str) -> float:
        """计算两个字符串的匹配度"""
        question = question.lower().strip()
        target = target.lower().strip()
        
        # 完全匹配
        if question == target:
            return 1.0
        
        # 包含关系
        if question in target or target in question:
            return 0.8
        
        # 单词重叠率
        q_words = set(question.split())
        t_words = set(target.split())
        
        if q_words and t_words:
            overlap = len(q_words & t_words) / len(q_words | t_words)
            return overlap
        
        return 0.0
    
    def save_knowledge_base(self, file_path: str = None) -> bool:
        """
        保存知识库到文件
        
        Args:
            file_path: 保存路径，如果为None则使用初始化时的file_path
            
        Returns:
            bool: 是否成功保存
        """
        if file_path is None:
            file_path = self.file_path
        
        if file_path is None or self.data is None:
            return False
        
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(self.data, f, ensure_ascii=False, indent=2)
            
            # 更新向量数据库
            if self.vector_db and self.dir_path:
                role = os.path.basename(file_path).replace('.json', '')
                self.vector_db.update_knowledge_base(role)
            
            return True
            
        except Exception as e:
            print(f"保存知识库文件 {file_path} 失败: {e}")
            return False
    
    def add_qa_pair(self, question: str, answer: str) -> bool:
        """
        添加问答对
        
        Args:
            question: 问题
            answer: 答案
            
        Returns:
            bool: 是否成功添加
        """
        if self.data is None:
            self.data = {
                'name': '未命名知识库',
                'description': '',
                'qa_pairs': []
            }
        
        if 'qa_pairs' not in self.data:
            self.data['qa_pairs'] = []
        
        self.data['qa_pairs'].append({
            'question': question,
            'answer': answer
        })
        
        # 更新向量数据库
        if self.vector_db and self.dir_path:
            role = None
            if self.file_path:
                role = os.path.basename(self.file_path).replace('.json', '')
            if role:
                self.vector_db.update_knowledge_base(role)
        
        return True
    
    def get_qa_pairs(self) -> List[Dict[str, str]]:
        """
        获取所有问答对
        
        Returns:
            List[Dict[str, str]]: 问答对列表
        """
        if self.data and 'qa_pairs' in self.data:
            return self.data['qa_pairs']
        return []
    
    def get_knowledge_base_info(self) -> Dict[str, Any]:
        """
        获取知识库信息
        
        Returns:
            Dict[str, Any]: 知识库信息
        """
        if self.data:
            return {
                'name': self.data.get('name', '未命名'),
                'description': self.data.get('description', ''),
                'qa_count': len(self.data.get('qa_pairs', [])),
                'file_path': self.file_path,
                'vector_db_enabled': self.vector_db is not None
            }
        elif self.knowledge_bases:
            return {
                'type': 'multi_kb',
                'kb_count': len(self.knowledge_bases),
                'kb_names': list(self.knowledge_bases.keys()),
                'vector_db_enabled': self.vector_db is not None
            }
        else:
            return {'type': 'empty'}

    @staticmethod
    def load_all_from_directory(dir_path: str) -> Dict[str, 'KnowledgeBase']:
        """
        从目录加载所有知识库文件（静态方法）
        
        Args:
            dir_path: 知识库目录路径
            
        Returns:
            Dict[str, KnowledgeBase]: 角色名到知识库实例的映射
        """
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
            return {}
        
        knowledge_bases = {}
        for filename in os.listdir(dir_path):
            if filename.endswith('.json'):
                kb_path = os.path.join(dir_path, filename)
                kb_name = filename.replace('.json', '')
                
                try:
                    kb = KnowledgeBase(file_path=kb_path)
                    if kb.data:
                        knowledge_bases[kb_name] = kb
                except Exception as e:
                    print(f"加载知识库 {filename} 失败: {e}")
        
        return knowledge_bases